use crate::ast::*;
use crate::error::{CompilerError, Result};
use inkwell::types::{BasicMetadataTypeEnum, BasicType, BasicTypeEnum};
use inkwell::values::FunctionValue;
use std::collections::HashMap;

use super::LLVMBackend;

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate LLVM IR for the entire NRC program.
    pub fn generate_program(&mut self, program: &Program) -> Result<()> {
        self.compute_overloaded_function_names(program);
        self.collect_function_declarations(program);

        // Add built-in functions
        self.add_builtin_functions()?;

        // Add extern declarations for imported functions and store their paths
        self.imported_module_paths = self.add_imported_function_declarations(program)?;

        // First pass: collect all struct/enum/trait declarations
        for declaration in &program.declarations {
            match declaration {
                Declaration::Struct(struct_decl) => {
                    if !struct_decl.generics.is_empty() {
                        continue;
                    }
                    self.struct_declarations
                        .insert(struct_decl.name.clone(), struct_decl.clone());
                }
                Declaration::Enum(enum_decl) => {
                    self.enum_declarations
                        .insert(enum_decl.name.clone(), enum_decl.clone());
                }
                Declaration::Type(type_decl) => {
                    // For Go-style structs defined with 'type', we need to create a StructDecl
                    if let Type::Struct(_struct_type) = &type_decl.type_def {
                        // Create a minimal StructDecl for lookup purposes
                        // The actual field information will be handled by the type system
                        let struct_decl = StructDecl {
                            name: type_decl.name.clone(),
                            generics: vec![],
                            fields: vec![], // Fields will be handled by the type system
                            methods: vec![],
                            location: type_decl.location,
                        };
                        self.struct_declarations
                            .insert(type_decl.name.clone(), struct_decl);
                    }
                }
                Declaration::Trait(trait_decl) => {
                    self.trait_declarations
                        .insert(trait_decl.name.clone(), trait_decl.clone());
                    self.ensure_trait_object_type(&trait_decl.name);
                    self.ensure_trait_rc_handle_type(&trait_decl.name);
                    self.ensure_trait_weak_handle_type(&trait_decl.name);
                }
                _ => {}
            }
        }

        // Second pass: process all impl blocks to add methods to struct declarations
        for declaration in &program.declarations {
            if let Declaration::Impl(impl_decl) = declaration {
                // Add methods to the corresponding struct declaration
                if let Some(struct_decl) = self.struct_declarations.get_mut(&impl_decl.type_name)
                {
                    // Clone the methods to avoid borrowing issues
                    let methods = impl_decl.methods.clone();
                    struct_decl.methods.extend(methods);
                }
            }
            if let Declaration::TraitImpl(trait_impl) = declaration {
                // Add methods from trait implementations to the corresponding struct declaration
                if let Some(struct_decl) = self.struct_declarations.get_mut(&trait_impl.for_type)
                {
                    // Clone the methods to avoid borrowing issues
                    let methods = trait_impl.methods.clone();
                    struct_decl.methods.extend(methods);
                }
                self.trait_impls.insert(
                    (trait_impl.trait_name.clone(), trait_impl.for_type.clone()),
                    trait_impl.clone(),
                );
            }
        }

        // Normalize trait method signatures using information from implementations when needed.
        // This backfills missing return types in trait declarations so that trait-object calls
        // can use the correct function signatures.
        for ((trait_name, _type_name), trait_impl) in &self.trait_impls {
            if let Some(trait_decl) = self.trait_declarations.get_mut(trait_name) {
                for impl_method in &trait_impl.methods {
                    if let Some(trait_method) =
                        trait_decl.methods.iter_mut().find(|m| m.name == impl_method.name)
                    {
                        if trait_method.return_type.is_none() && impl_method.return_type.is_some()
                        {
                            trait_method.return_type = impl_method.return_type.clone();
                        }
                    }
                }
            }
        }

        // Third pass: declare all functions (without generating bodies)
        for declaration in &program.declarations {
            self.declare_functions(declaration)?;
        }

        // Fourth pass: generate enum constants
        for declaration in &program.declarations {
            if let Declaration::Enum(enum_decl) = declaration {
                self.generate_enum_constants(enum_decl)?;
            }
        }

        // Fifth pass: generate function bodies
        for declaration in &program.declarations {
            self.generate_function_bodies(declaration)?;
        }

        Ok(())
    }

    fn compute_overloaded_function_names(&mut self, program: &Program) {
        let mut counts: HashMap<String, usize> = HashMap::new();
        for declaration in &program.declarations {
            if let Declaration::Function(func) = declaration {
                *counts.entry(func.name.clone()).or_insert(0) += 1;
            }
        }
        self.overloaded_function_names = counts
            .into_iter()
            .filter(|(_, count)| *count > 1)
            .map(|(name, _)| name)
            .collect();
    }

    fn collect_function_declarations(&mut self, program: &Program) {
        self.function_decls.clear();
        for declaration in &program.declarations {
            if let Declaration::Function(func) = declaration {
                if !self.overloaded_function_names.contains(&func.name) {
                    self.function_decls.insert(func.name.clone(), func.clone());
                }
            }
        }
    }

    /// Declare all functions (first pass - no body generation)
    fn declare_functions(&mut self, declaration: &Declaration) -> Result<()> {
        match declaration {
            Declaration::Function(func) => {
                self.declare_function(func)?;
            }
            Declaration::Struct(struct_decl) => {
                // Declare methods for Go-style structs
                for method in &struct_decl.methods {
                    self.declare_function(method)?;
                }
            }
            Declaration::Impl(impl_decl) => {
                // Declare methods for Rust-style impl blocks
                for method in &impl_decl.methods {
                    let method_name = format!("{}.{}", impl_decl.type_name, method.name);

                    // Skip if already declared
                    let mut method_copy = method.clone();
                    method_copy.name = method_name.clone();
                    let symbol_name = self.function_symbol_name(&method_copy);
                    if self.function_map.contains_key(&symbol_name) {
                        continue;
                    }

                    self.declare_function(&method_copy)?;
                }
            }
            Declaration::TraitImpl(trait_impl) => {
                // Declare methods for trait implementations
                for method in &trait_impl.methods {
                    let method_name = format!("{}.{}", trait_impl.for_type, method.name);

                    // Skip if already declared
                    let mut method_copy = method.clone();
                    method_copy.name = method_name.clone();
                    let symbol_name = self.function_symbol_name(&method_copy);
                    if self.function_map.contains_key(&symbol_name) {
                        continue;
                    }

                    self.declare_function(&method_copy)?;
                }
            }
            _ => {
                // Other declaration types don't need function declarations
            }
        }
        Ok(())
    }

    /// Generate function bodies (second pass)
    fn generate_function_bodies(&mut self, declaration: &Declaration) -> Result<()> {
        match declaration {
            Declaration::Function(func) => {
                if func.generics.is_empty() {
                    if let Some(ref body) = func.body {
                        let symbol_name = self.function_symbol_name(func);
                        let function_value =
                            *self.function_map.get(&symbol_name).ok_or_else(|| {
                                CompilerError::internal(&format!(
                                    "Function {} not declared",
                                    func.name
                                ))
                            })?;
                        self.current_function = Some(function_value);
                        let previous_return = self.current_function_return_type.take();
                        self.current_function_return_type = func.return_type.clone();
                        let body_result =
                            self.generate_function_body(function_value, &func.parameters, body);
                        self.current_function_return_type = previous_return;
                        body_result?;
                    }
                }
            }
            Declaration::Struct(struct_decl) => {
                // Generate method bodies for Go-style structs
                for method in &struct_decl.methods {
                    if !method.generics.is_empty() {
                        continue;
                    }
                    if let Some(ref body) = method.body {
                        let symbol_name = self.function_symbol_name(method);
                        let function_value =
                            *self.function_map.get(&symbol_name).ok_or_else(|| {
                                CompilerError::internal(&format!(
                                    "Method {} not declared",
                                    method.name
                                ))
                            })?;
                        self.current_function = Some(function_value);
                        let previous_return = self.current_function_return_type.take();
                        self.current_function_return_type = method.return_type.clone();
                        let body_result =
                            self.generate_function_body(function_value, &method.parameters, body);
                        self.current_function_return_type = previous_return;
                        body_result?;
                    }
                }
            }
            Declaration::Impl(impl_decl) => {
                // Generate method bodies for Rust-style impl blocks
                for method in &impl_decl.methods {
                    if !method.generics.is_empty() {
                        continue;
                    }
                    if let Some(ref body) = method.body {
                        // Use fully qualified method name to avoid conflicts
                        let method_name = format!("{}.{}", impl_decl.type_name, method.name);
                        let mut method_copy = method.clone();
                        method_copy.name = method_name.clone();
                        let symbol_name = self.function_symbol_name(&method_copy);
                        let function_value =
                            *self.function_map.get(&symbol_name).ok_or_else(|| {
                                CompilerError::internal(&format!(
                                    "Method {} not declared",
                                    method_name
                                ))
                            })?;
                        self.current_function = Some(function_value);
                        let previous_return = self.current_function_return_type.take();
                        self.current_function_return_type = method.return_type.clone();
                        let body_result =
                            self.generate_function_body(function_value, &method.parameters, body);
                        self.current_function_return_type = previous_return;
                        body_result?;
                    }
                }
            }
            Declaration::TraitImpl(trait_impl) => {
                // Generate method bodies for trait implementations
                for method in &trait_impl.methods {
                    if !method.generics.is_empty() {
                        continue;
                    }
                    if let Some(ref body) = method.body {
                        // Use fully qualified method name to avoid conflicts
                        let method_name = format!("{}.{}", trait_impl.for_type, method.name);
                        let mut method_copy = method.clone();
                        method_copy.name = method_name.clone();
                        let symbol_name = self.function_symbol_name(&method_copy);
                        let function_value =
                            *self.function_map.get(&symbol_name).ok_or_else(|| {
                                CompilerError::internal(&format!(
                                    "Method {} not declared",
                                    method_name
                                ))
                            })?;
                        self.current_function = Some(function_value);
                        let previous_return = self.current_function_return_type.take();
                        self.current_function_return_type = method.return_type.clone();
                        let body_result =
                            self.generate_function_body(function_value, &method.parameters, body);
                        self.current_function_return_type = previous_return;
                        body_result?;
                    }
                }
            }
            _ => {
                // Other declaration types don't have function bodies
            }
        }
        Ok(())
    }

    /// Declare a function (signature only, no body)
    fn declare_function(&mut self, func: &FunctionDecl) -> Result<()> {
        if !func.generics.is_empty() {
            // Generic templates are instantiated separately; skip declaring them directly.
            return Ok(());
        }
        // Skip if function already declared
        let symbol_name = self.function_symbol_name(func);
        if self.function_map.contains_key(&symbol_name) {
            return Ok(());
        }

        // Get parameter types
        let param_types: Vec<BasicTypeEnum> = func
            .parameters
            .iter()
            .map(|param| self.nrc_type_to_llvm_type(&Some(param.param_type.clone())))
            .collect::<Result<Vec<_>>>()?;

        // Convert to metadata types
        let metadata_types: Vec<BasicMetadataTypeEnum> =
            param_types.iter().map(|t| (*t).into()).collect();

        // Create function type
        let function = if func.return_type.is_none() {
            let function_type = self.context.void_type().fn_type(&metadata_types, false);
            self.module.add_function(&symbol_name, function_type, None)
        } else {
            let return_type = self.nrc_type_to_llvm_type(&func.return_type)?;
            let function_type = return_type.fn_type(&metadata_types, false);
            self.module.add_function(&symbol_name, function_type, None)
        };

        // Set parameter names
        for (i, param) in func.parameters.iter().enumerate() {
            if let Some(llvm_param) = function.get_nth_param(i as u32) {
                llvm_param.set_name(&param.name);
            }
        }

        // Store function in map
        self.function_map.insert(symbol_name, function);

        Ok(())
    }

    /// Generate LLVM IR for a function
    // fn generate_function(&mut self, func: &FunctionDecl) -> Result<()> {

    //     // Get parameter types
    //     let param_types: Vec<BasicTypeEnum> = func.parameters
    //         .iter()
    //         .map(|param| self.nrc_type_to_llvm_type(&Some(param.param_type.clone())))
    //         .collect::<Result<Vec<_>>>()?;

    //     // Convert to metadata types
    //     let metadata_types: Vec<inkwell::types::BasicMetadataTypeEnum> = param_types
    //         .iter()
    //         .map(|t| (*t).into())
    //         .collect();

    //     // Create function type - handle void return type specially
    //     let function = if func.return_type.is_none() {
    //         // Void function
    //         let function_type = self.context.void_type().fn_type(&metadata_types, false);
    //         self.module.add_function(&func.name, function_type, None)
    //     } else {
    //         // Function with return value
    //         let return_type = self.nrc_type_to_llvm_type(&func.return_type)?;
    //         let function_type = return_type.fn_type(&metadata_types, false);
    //         self.module.add_function(&func.name, function_type, None)
    //     };

    //     // Set parameter names
    //     for (i, param) in func.parameters.iter().enumerate() {
    //         if let Some(llvm_param) = function.get_nth_param(i as u32) {
    //             llvm_param.set_name(&param.name);
    //         }
    //     }

    //     // Store function in map
    //     self.function_map.insert(func.name.clone(), function);

    //     // Generate function body if it exists
    //     if let Some(ref body) = func.body {
    //         self.current_function = Some(function);
    //         self.generate_function_body(function, &func.parameters, body)?;
    //     }

    //     Ok(())
    // }

    /// Generate function body
    fn generate_function_body(
        &mut self,
        function: FunctionValue<'ctx>,
        parameters: &[crate::ast::types::Parameter],
        body: &Block,
    ) -> Result<()> {
        // Create entry block
        let entry_block = self.context.append_basic_block(function, "entry");
        self.builder.position_at_end(entry_block);

        // 清空作用域栈和stack slices列表
        self.scope_stack.clear();
        self.stack_slices.clear();

        // 创建函数作用域
        self.enter_scope();

        // 为函数参数创建alloca并存储值（就像局部变量一样）
        // 这样所有变量都通过alloca访问，逻辑更一致
        for (i, param_decl) in parameters.iter().enumerate() {
            if let Some(llvm_param) = function.get_nth_param(i as u32) {
                // 获取参数的LLVM类型
                let param_llvm_type =
                    self.nrc_type_to_llvm_type(&Some(param_decl.param_type.clone()))?;

                // 创建alloca
                let alloca = self
                    .builder
                    .build_alloca(param_llvm_type, &param_decl.name)?;

                // 将参数值存储到alloca中
                let _ = self.builder.build_store(alloca, llvm_param);

                // 将alloca添加到变量映射中
                self.variable_map
                    .insert(param_decl.name.clone(), alloca.into());

                // 同时将参数类型添加到类型映射中
                self.variable_type_map
                    .insert(param_decl.name.clone(), param_decl.param_type.clone());
            }
        }

        // 检查是否是单个表达式语句（隐式返回）
        let is_single_expression =
            body.statements.len() == 1 && matches!(body.statements[0], Statement::Expression(_));

        // Generate statements
        let mut has_return = false;
        for stmt in &body.statements {
            if matches!(stmt, Statement::Return(_)) {
                has_return = true;
                // 遇到return语句，先执行所有defer，然后清理stack slices，最后生成return
                self.execute_all_defer_statements()?;
                self.cleanup_stack_slices()?;
                self.generate_statement(stmt)?;
                break; // 停止生成后续语句
            } else if is_single_expression && function.get_type().get_return_type().is_some() {
                // 单个表达式语句作为隐式返回
                if let Statement::Expression(expr) = stmt {
                    let expr_value = self.generate_expression(expr)?;
                    self.execute_all_defer_statements()?;
                    self.cleanup_stack_slices()?;
                    self.cleanup_manual_allocated_variables()?;
                    let _ = self.builder.build_return(Some(&expr_value));
                    has_return = true;
                    break;
                }
            } else {
                self.generate_statement(stmt)?;
            }
        }

        // 如果没有遇到return语句，在函数结束前执行所有defer语句（LIFO顺序）和清理stack slices
        if !has_return {
            // Create a function end block where after_recover blocks will jump to
            let function_end_block = self.context.append_basic_block(function, "function_end");

            // Branch from current position to function_end_block
            self.builder
                .build_unconditional_branch(function_end_block)?;
            self.builder.position_at_end(function_end_block);

            // Connect all recover after_recover blocks to function_end_block
            // This ensures that when panic is caught, execution continues at function end
            for after_recover_block in &self.recover_after_blocks {
                // Check if after_recover_block has no terminator (empty block)
                if after_recover_block.get_terminator().is_none() {
                    let saved_position = self.builder.get_insert_block();
                    self.builder.position_at_end(*after_recover_block);
                    // Branch to function_end_block
                    self.builder
                        .build_unconditional_branch(function_end_block)?;
                    if let Some(saved) = saved_position {
                        self.builder.position_at_end(saved);
                    }
                }
            }

            self.execute_all_defer_statements()?;
            self.cleanup_stack_slices()?;

            // 清理所有手动分配的变量
            self.cleanup_manual_allocated_variables()?;

            // Add return statement if function returns void or没有显式返回
            if function.get_type().get_return_type().is_none() {
                let _ = self.builder.build_return(None);
            } else {
                // 如果函数有返回值但没有显式返回，添加默认返回
                let return_type = function.get_type().get_return_type().unwrap();
                match return_type {
                    BasicTypeEnum::IntType(int_type) => {
                        let zero = int_type.const_int(0, false);
                        let _ = self.builder.build_return(Some(&zero));
                    }
                    _ => {
                        let _ = self.builder.build_return(None);
                    }
                }
            }
        } else {
            // If there's a return statement, connect recover after_recover blocks to function end
            // But since we already returned, these blocks should just contain return statements
            // For now, leave them empty (they won't be reached)
        }

        // 清空变量映射和类型映射，防止跨函数引用
        self.variable_map.clear();
        self.variable_type_map.clear();

        Ok(())
    }}
